import os
import os.path as osp

import cv2
import numpy as np
from tqdm import tqdm


if __name__ == "__main__":
    with open("cameraMatrix.txt", 'r') as f:
        data = f.read().splitlines()

    cameraMatrix = np.array(eval(data[0]), dtype=np.float32).reshape(3, 3)
    distCoeffs = np.array(eval(data[1]), dtype=np.float32)

    image = cv2.imread("40.jpg")
    imgSize = (image.shape[1], image.shape[0])



    # 计算新的相机参数以及ROI: 需要保留原图的整图信息时
    new_cameraMatrix, roi = cv2.getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, imgSize, 1, imgSize)


    # 计算畸变纠正map
    map1, map2 = cv2.initUndistortRectifyMap(cameraMatrix, distCoeffs, None, new_cameraMatrix, imgSize, cv2.CV_16SC2)
    undistort_image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0))

    # 若单张图片，可以使用cv2.undistort()方法进行纠正
    # undistort_image = cv2.undistort(image, cameraMatrix, distCoeffs, None, new_cameraMatrix)


    # show
    # cv2.rectangle(undistort_image, (roi[0], roi[1]), (roi[0] + roi[2], roi[1] + roi[3]), (0, 0, 255), 2)
    print(f"image shape: {image.shape[:2]}")
    print(f"undistort_image shape: {undistort_image.shape[:2]}")
    # print(f"result_image shape: {result_image.shape[:2]}")
    cv2.imshow("image", image)
    cv2.imshow("undistort_image", undistort_image)
    # cv2.imshow("result_image", result_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    cv2.imwrite("undistort_image.jpg", undistort_image)
    pass